package outparser

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

// OutputFormat represents the format of the output
type OutputFormat string

const (
	FormatJSON     OutputFormat = "json"
	FormatYAML     OutputFormat = "yaml"
	FormatMarkdown OutputFormat = "markdown"
	FormatText     OutputFormat = "text"
	FormatCode     OutputFormat = "code"
)

// OutputParser defines the interface for parsing LLM outputs
type OutputParser interface {
	// Parse parses the output into a structured format
	Parse(output string) (interface{}, error)

	// Format returns the expected format for the output
	Format() OutputFormat

	// Example returns an example of the expected output format
	Example() string
}

// JSONParser parses JSON output
type JSONParser struct {
	schema map[string]interface{} // expected JSON schema
}

// NewJSONParser creates a new JSON parser
func NewJSONParser(schema map[string]interface{}) *JSONParser {
	return &JSONParser{schema: schema}
}

func (p *JSONParser) Parse(output string) (interface{}, error) {
	// Extract JSON from the output
	jsonStr := extractJSON(output)
	if jsonStr == "" {
		return nil, fmt.Errorf("no JSON found in output")
	}

	var result interface{}
	if err := json.Unmarshal([]byte(jsonStr), &result); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %v", err)
	}

	return result, nil
}

func (p *JSONParser) Format() OutputFormat {
	return FormatJSON
}

func (p *JSONParser) Example() string {
	example, _ := json.MarshalIndent(p.schema, "", "  ")
	return string(example)
}

// CodeParser parses code output
type CodeParser struct {
	language string
}

// NewCodeParser creates a new code parser
func NewCodeParser(language string) *CodeParser {
	return &CodeParser{language: language}
}

func (p *CodeParser) Parse(output string) (interface{}, error) {
	// Extract code blocks
	code := extractCodeBlock(output, p.language)
	if code == "" {
		return nil, fmt.Errorf("no code block found for language: %s", p.language)
	}

	return map[string]interface{}{
		"language": p.language,
		"code":     code,
	}, nil
}

func (p *CodeParser) Format() OutputFormat {
	return FormatCode
}

func (p *CodeParser) Example() string {
	return fmt.Sprintf("'''%s\n// Your code here\n'''", p.language)
}

// ListParser parses list output
type ListParser struct {
	itemPrefix string
}

// NewListParser creates a new list parser
func NewListParser(itemPrefix string) *ListParser {
	return &ListParser{itemPrefix: itemPrefix}
}

func (p *ListParser) Parse(output string) (interface{}, error) {
	lines := strings.Split(output, "\n")
	var items []string

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, p.itemPrefix) {
			item := strings.TrimPrefix(line, p.itemPrefix)
			items = append(items, strings.TrimSpace(item))
		}
	}

	if len(items) == 0 {
		return nil, fmt.Errorf("no list items found with prefix: %s", p.itemPrefix)
	}

	return items, nil
}

func (p *ListParser) Format() OutputFormat {
	return FormatText
}

func (p *ListParser) Example() string {
	return fmt.Sprintf("%s First item\n%s Second item\n%s Third item",
		p.itemPrefix, p.itemPrefix, p.itemPrefix)
}

// Helper functions

func extractJSON(text string) string {
	// Find text between { and } with the most balanced braces
	re := regexp.MustCompile(`{(?:[^{}]|{(?:[^{}]|{[^{}]*})*})*}`)
	if match := re.FindString(text); match != "" {
		return match
	}
	return ""
}

func extractCodeBlock(text, language string) string {
	// Find code blocks with the specified language
	pattern := fmt.Sprintf("(?s)'''%s\\n(.+?)'''", language)
	re := regexp.MustCompile(pattern)
	if matches := re.FindStringSubmatch(text); len(matches) > 1 {
		return strings.TrimSpace(matches[1])
	}
	return ""
}

// Common output schemas
var (
	CodeAnalysisSchema = map[string]interface{}{
		"overview": "string",
		"issues": []map[string]interface{}{
			{
				"type":        "string",
				"severity":    "string",
				"description": "string",
				"line":        "number",
			},
		},
		"suggestions": []string{},
		"metrics": map[string]interface{}{
			"complexity": "number",
			"loc":        "number",
		},
	}

	TaskResultSchema = map[string]interface{}{
		"success": "boolean",
		"result":  "string",
		"error":   "string",
		"metrics": map[string]interface{}{
			"duration": "number",
			"memory":   "number",
		},
	}
)
